{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A reimplementation of \"Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains\" (https://github.com/ndahlquist/pytorch-fourier-feature-networks/tree/master). We are going to represent a 2D image with an MLP and verify the effectiveness of Fourier Features. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision.io import read_image, ImageReadMode\n",
    "from torchvision.transforms.functional import to_pil_image\n",
    "\n",
    "from tqdm import tqdm\n",
    "from einops import rearrange\n",
    "\n",
    "def viz_image(pt_img: torch.Tensor):\n",
    "    pil_img = to_pil_image(pt_img)\n",
    "    display(pil_img)\n",
    "\n",
    "    \n",
    "input_image = read_image('misuzu.png', ImageReadMode.RGB)\n",
    "input_image = input_image.to(torch.float32) / 255\n",
    "input_image = input_image.unsqueeze(0)\n",
    "input_image = F.interpolate(input_image, (256, 256), mode='bilinear')\n",
    "viz_image(input_image[0])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self, in_c, out_c=3, hiden_states=256):\n",
    "        super().__init__()\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Conv2d(in_c, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),\n",
    "            nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),\n",
    "            nn.Conv2d(hiden_states, hiden_states, 1), nn.ReLU(), nn.BatchNorm2d(hiden_states),\n",
    "            nn.Conv2d(hiden_states, out_c, 1), nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.mlp(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H, W = input_image.shape[2:]\n",
    "\n",
    "h_coord = torch.linspace(0, 1, H)\n",
    "w_coord = torch.linspace(0, 1, W)\n",
    "grid = torch.stack(torch.meshgrid([h_coord, w_coord]), -1).permute(2, 0, 1).unsqueeze(0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "model = MLP(2).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
    "n_loops = 400\n",
    "input_image = input_image.to(device)\n",
    "grid = grid.to(device)\n",
    "for epoch in tqdm(range(n_loops)):\n",
    "    output = model(grid)\n",
    "    loss = F.l1_loss(output, input_image)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if epoch % 100 == 0 or epoch == n_loops - 1:\n",
    "        viz_image(output[0])\n",
    "        print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FourierFeature(nn.Module):\n",
    "    def __init__(self, in_c, out_c, scale):\n",
    "        super().__init__()\n",
    "        fourier_basis = torch.randn(in_c, out_c // 2) * scale\n",
    "        self.register_buffer('_fourier_basis', fourier_basis)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        N, C, H, W = x.shape\n",
    "        x = rearrange(x, 'n c h w -> (n h w) c')\n",
    "        x = x @ self._fourier_basis\n",
    "        x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)\n",
    "            \n",
    "        x = 2 * torch.pi * x\n",
    "        x = torch.cat([torch.sin(x), torch.cos(x)], dim=1) \n",
    "        return x\n",
    "        \n",
    "feature_length = 256\n",
    "model = MLP(feature_length).to(device)\n",
    "fourier_feature = FourierFeature(2, feature_length, 10).to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
    "n_loops = 400\n",
    "for epoch in tqdm(range(n_loops)):\n",
    "    x = fourier_feature(grid)\n",
    "    output = model(x)\n",
    "    loss = F.l1_loss(output, input_image)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if epoch % 100 == 0 or epoch == n_loops - 1:\n",
    "        viz_image(output[0])\n",
    "        print(loss.item())\n",
    "prev_output = output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N, C, H, W = grid.shape\n",
    "tx = 50 / H\n",
    "ty = 0\n",
    "theta = torch.tensor(torch.pi * 1 / 8)\n",
    "affine_matrix = torch.tensor([\n",
    "    [torch.cos(theta), -torch.sin(theta), tx],\n",
    "    [torch.sin(theta), torch.cos(theta), ty],\n",
    "    [0, 0, 1]\n",
    "]\n",
    ").to(device)\n",
    "grid_ext = torch.ones(N, 3, H, W).to(device)\n",
    "grid_ext[:, :2] = grid.clone()\n",
    "grid_ext = grid_ext.permute(0, 2, 3, 1)\n",
    "grid_ext = (grid_ext @ affine_matrix.T)\n",
    "grid_ext = grid_ext.permute(0, 3, 1, 2)[:, :2]\n",
    "\n",
    "x = fourier_feature(grid_ext)\n",
    "output = model(x)\n",
    "viz_image(output[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def aff_transform(model, x, aff):\n",
    "    N, C, H, W = x.shape\n",
    "    x = rearrange(x, 'n c h w -> (n h w) c')\n",
    "    \n",
    "    if aff is not None:\n",
    "        phases = torch.zeros(1, model._fourier_basis.shape[1]).to(x.device)\n",
    "        phases = phases + (model._fourier_basis.T @ aff[:2, 2:]).T\n",
    "        freq = model._fourier_basis.T @ aff[:2, :2]\n",
    "        x = x @ freq.T\n",
    "        x = x + phases\n",
    "    else:\n",
    "        x = x @ model._fourier_basis\n",
    "    x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)\n",
    "    \n",
    "    \n",
    "    x = 2 * torch.pi * x\n",
    "    x = torch.cat([torch.sin(x), torch.cos(x)], dim=1) \n",
    "    return x\n",
    "\n",
    "x = aff_transform(fourier_feature, grid, affine_matrix)\n",
    "output2 = model(x)\n",
    "viz_image(output2[0])\n",
    "print(F.l1_loss(output2, output).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FourierFeature(nn.Module):\n",
    "    def __init__(self, in_c, out_c, scale):\n",
    "        super().__init__()\n",
    "        fourier_basis = torch.randn(in_c, out_c) * scale\n",
    "        self.register_buffer('_fourier_basis', fourier_basis)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        N, C, H, W = x.shape\n",
    "        x = rearrange(x, 'n c h w -> (n h w) c')\n",
    "        x = x @ self._fourier_basis\n",
    "        x = rearrange(x, '(n h w) c -> n c h w', h = H, w = W)\n",
    "            \n",
    "        x = 2 * torch.pi * x\n",
    "        x = torch.sin(x)\n",
    "        return x\n",
    "\n",
    "feature_length = 256\n",
    "model2 = MLP(feature_length).to(device)\n",
    "fourier_feature = FourierFeature(2, feature_length, 10).to(device)\n",
    "optimizer = torch.optim.Adam(model2.parameters(), lr=1e-4)\n",
    "n_loops = 500\n",
    "for epoch in tqdm(range(n_loops)):\n",
    "    x = fourier_feature(grid + 10)\n",
    "    output = model2(x)\n",
    "    loss = F.l1_loss(output, input_image)\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    \n",
    "    if epoch % 100 == 0 or epoch == n_loops - 1:\n",
    "        viz_image(output[0])\n",
    "        print(loss.item())\n",
    "        print(F.l1_loss(output, prev_output).item())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pt",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
